from pyfeasc.paasc import *

# 测试
if __name__ == '__main__':
	# 设置随机种子保证可重复性
	np.random.seed(42)
	
	# 生成相关数据
	n = 30
	x = np.random.normal(size=n)
	y = 0.5 * x + np.random.normal(scale=0.5, size=n)  # y 与 x相关
	z = pd.DataFrame({
			'z1': 0.3 * x + 0.3 * y + np.random.normal(scale=0.5, size=n),  # z1 与 x, y 都相关
			'z2': np.random.normal(size=n),  # z2 与 x, y 无关
			'z3': 0.5 * x + np.random.normal(scale=0.5, size=n)  # z3 只与 x 相关
	})
	
	# 生成空间坐标数据
	spatial_x = np.random.uniform(size=n)
	spatial_y = np.random.uniform(size=n)
	spatial_z = pd.DataFrame({
			'sz1': 0.5 * spatial_x + 0.5 * spatial_y + np.random.normal(scale=0.2, size=n),  # 与空间位置强相关
			'sz2': np.random.normal(size=n),  # 与空间位置无关
			'sz3': 0.3 * spatial_x + np.random.normal(scale=0.5, size=n)  # 只与x坐标相关
	})
	
	# 测试 y 与 x 的互信息
	mi_result = calculate_mi(x, y)
	print("MI between x and y:")
	print(mi_result)
	
	# 测试 z 与 x 的互信息
	mi_result_z = calculate_mi(x, z)
	print("\nMI between x and z:")
	print(mi_result_z)
	
	# 测试空间互信息
	spatial_mi_result = calculate_spatial_mi(spatial_x, spatial_y, spatial_z)
	print("\nSpatial MI:")
	print(spatial_mi_result)
